"""代码生成测试辅助工具。

本模块提供了用于测试代码生成功能的辅助函数
"""

import textwrap
import inspect
from pathlib import Path
from typing import Union, Callable, Optional

from mcpy.compiler import MCPyCompiler


def read_lines(code: str):
    """读取代码行，去除空行和首尾空白，并规范化行内空格"""
    import re

    lines = []
    for line in code.splitlines():
        stripped = line.strip()
        if stripped:
            # 将多个连续空格替换为单个空格，处理clang-format格式化的差异
            normalized = re.sub(r"\s+", " ", stripped)
            lines.append(normalized)
    return lines


def normalize_code(code: str) -> str:
    return "\n".join(read_lines(code))


def compare_code(code: str, expected: str) -> bool:
    """按行迭代比较两个代码片段，忽略空行和首尾空白。

    Args:
        code: 生成的完整代码
        expected: 期望包含的代码片段

    Returns:
        bool: 如果期望的代码片段在生成的代码中找到则返回True
    """

    # 将期望的代码片段转换为列表用于比较
    expected_lines = read_lines(expected)
    code_lines = read_lines(code)
    if not expected_lines:
        return True

    if len(expected_lines) > len(code_lines):
        return False

    for i in range(len(code_lines) - len(expected_lines) + 1):
        if code_lines[i] == expected_lines[0]:
            if code_lines[i + 1 : i + len(expected_lines)] == expected_lines[1:]:
                return True

        i += 1

    return False


def expand_regex_markers(pattern: str) -> str:
    """将 ${} 标记的正则部分展开，其他部分转义为字面量

    ${} 内可以是：
    1. 简写形式（自动转换为正则）
    2. 完整正则表达式（直接使用）
    3. 任意自定义正则

    Args:
        pattern: 包含 ${} 标记的模式字符串

    Returns:
        完整的正则表达式字符串

    示例:
        >>> # 使用简写
        >>> expand_regex_markers("struct __lambda_${w+}_${d+}_${d+}")
        'struct __lambda_\\w+_\\d+_\\d+'

        >>> # 使用完整正则
        >>> expand_regex_markers("value: ${\\d+\\.\\d+}")
        'value: \\d+\\.\\d+'

        >>> # 混合使用
        >>> expand_regex_markers("mc::variant<${.+?}>(${.*?})")
        'mc::variant<.+?>\\(.*?\\)'

    支持的简写（可扩展）：
        ${w+}  ${w*}  ${w}   - 标识符相关
        ${d+}  ${d*}  ${d}   - 数字相关
        ${s+}  ${s*}  ${s}   - 空白相关
        ${.*?} ${.+?}        - 非贪婪匹配
        ${.*}  ${.+}         - 贪婪匹配

    高级用法：
        ${\w+_\d+}           - 直接写完整正则
        ${int|bool|float}    - 匹配多个选项
        ${[a-zA-Z_]\w*}      - 自定义字符类
    """
    import re

    result = []
    i = 0

    while i < len(pattern):
        # 查找下一个 ${ 标记
        start = pattern.find("${", i)

        if start == -1:
            # 没有更多标记，转义剩余部分
            result.append(re.escape(pattern[i:]))
            break

        # 转义标记之前的字面量部分
        if start > i:
            result.append(re.escape(pattern[i:start]))

        # 查找对应的 }
        end = pattern.find("}", start + 2)
        if end == -1:
            # 未找到闭合的 }，当作字面量处理
            result.append(re.escape(pattern[start:]))
            break

        # 提取 ${} 内的正则表达式
        regex_content = pattern[start + 2 : end]

        # 智能处理简写：只在完全匹配时转换，否则当作原始正则
        # 这样既支持简写，也支持完整正则表达式
        shorthand_map = {
            "w+": r"\w+",
            "w*": r"\w*",
            "w": r"\w",
            "d+": r"\d+",
            "d*": r"\d*",
            "d": r"\d",
            "s+": r"\s+",
            "s*": r"\s*",
            "s": r"\s",
            ".*?": r".*?",
            ".+?": r".+?",  # 非贪婪匹配
            ".*": r".*",
            ".+": r".+",  # 贪婪匹配
        }

        # 如果是简写形式，转换；否则保持原样（作为完整正则）
        if regex_content in shorthand_map:
            regex_content = shorthand_map[regex_content]
        # 如果包含反斜杠，说明已经是转义的正则，保持原样
        # 否则按原样使用（用户可以写任意正则表达式）

        result.append(regex_content)
        i = end + 1

    return "".join(result)


def compare_code_with_regex(
    code: str, expected_pattern: str, use_markers: bool = False
) -> bool:
    """使用正则表达式比较代码片段

    支持在期望的代码中使用正则表达式来匹配变化的部分（如行号、变量名）。
    每行的正则模式会逐行匹配代码。

    自动忽略首尾空白字符，保持测试代码的可读性。

    Args:
        code: 生成的完整代码
        expected_pattern: 期望的代码模式（每行可以是正则表达式）
        use_markers: 是否使用 ${} 标记（推荐，可读性高）

    Returns:
        bool: 如果找到匹配的代码块则返回 True

    示例:
        >>> # 使用 ${} 标记（推荐，可读性高，支持缩进）
        >>> pattern = '''
        ...     struct __lambda_${w+}_${d+}_${d+} {
        ...         auto operator()(int x) const {
        ...             return x * x;
        ...         }
        ...     };
        ... '''
        >>> compare_code_with_regex(code, pattern, use_markers=True)

        >>> # 原始正则（需要转义）
        >>> pattern = r'''
        ... mc::variant square = mc::variant\(mc::make_shared<__lambda_\w+_\d+_\d+>\(\)\);
        ... '''
        >>> compare_code_with_regex(code, pattern)
    """
    import re

    # 规范化代码行（去掉首尾空白）
    code_lines = read_lines(code)

    # 如果使用标记，先规范化每一行，然后逐行展开标记
    if use_markers:
        # 先按行分割并去掉空白，然后逐行展开标记
        pattern_lines_raw = read_lines(expected_pattern)
        pattern_lines = []
        for line in pattern_lines_raw:
            # 对每一行展开标记
            expanded = expand_regex_markers(line)
            pattern_lines.append(expanded)
    else:
        # 原始正则模式，直接规范化
        pattern_lines = read_lines(expected_pattern)

    if not pattern_lines:
        return True

    if len(pattern_lines) > len(code_lines):
        return False

    # 编译所有正则表达式模式
    compiled_patterns = []
    for pattern_line in pattern_lines:
        try:
            compiled_patterns.append(re.compile(pattern_line))
        except re.error:
            # 如果不是有效的正则表达式，当作普通字符串
            compiled_patterns.append(None)

    # 滑动窗口匹配
    for i in range(len(code_lines) - len(pattern_lines) + 1):
        # 尝试从位置 i 开始匹配所有模式行
        all_match = True
        for j, pattern in enumerate(compiled_patterns):
            code_line = code_lines[i + j]
            pattern_line = pattern_lines[j]

            if pattern is not None:
                # 使用正则匹配
                if not pattern.search(code_line):
                    all_match = False
                    break
            else:
                # 使用精确匹配
                if code_line != pattern_line:
                    all_match = False
                    break

        if all_match:
            return True

    return False


def assert_code_contains(
    code: str,
    expected: str,
    message: str = "",
    use_regex: bool = False,
    use_markers: bool = False,
) -> None:
    """断言完整代码包含期望的代码片段。

    使用逐行比较来处理换行和空白字符的不一致问题，能更好地处理模板中的连续换行。

    支持三种匹配模式：
    1. 普通匹配：精确匹配代码
    2. ${} 标记匹配：使用 ${} 标记正则部分（推荐，可读性高）
    3. 原始正则匹配：完整正则表达式（需要转义）

    Args:
        code: 完整代码
        expected: 期望包含的代码片段
        message: 断言失败时的错误消息
        use_regex: 是否使用原始正则表达式（默认 False）
        use_markers: 是否使用 ${} 标记（默认 False，推荐）

    ${} 标记语法：
        ${w+}     - 匹配标识符（\w+）
        ${d+}     - 匹配数字（\d+）
        ${.*?}    - 匹配任意内容（非贪婪）
        ${.+?}    - 匹配任意非空内容（非贪婪）
        ${\\w+}   - 完整正则：\w+
        ${type|int|bool} - 匹配多个选项之一

    示例:
        >>> # 普通匹配
        >>> assert_code_contains(code, "int foo() { return 42; }")

        >>> # ${} 标记匹配（推荐，可读性最高）
        >>> expected = '''
        ... struct __lambda_${w+}_${d+}_${d+} {
        ...     auto operator()(int x) const {
        ...         return x * x;
        ...     }
        ... };
        ... '''
        >>> assert_code_contains(code, expected, use_markers=True)

        >>> # 原始正则匹配（需要转义）
        >>> expected = r'''
        ... mc::variant square = mc::variant\(mc::make_shared<__lambda_\w+_\d+_\d+>\(\)\);
        ... '''
        >>> assert_code_contains(code, expected, use_regex=True)
    """
    if use_markers or use_regex:
        if not compare_code_with_regex(code, expected, use_markers=use_markers):
            error_msg = "代码片段不匹配"
            if use_markers:
                error_msg += "（${} 标记）"
            elif use_regex:
                error_msg += "（正则）"

            if message:
                error_msg = f"{message}: {error_msg}"

            if use_markers:
                display_pattern = f"原始模式:\n{expected}"
            else:
                display_pattern = expected

            error_msg += f"\n\n{display_pattern}\n\n实际代码:\n{normalize_code(code)}"
            raise AssertionError(error_msg)
    else:
        if not compare_code(code, expected):
            error_msg = "代码片段不匹配"
            if message:
                error_msg = f"{message}: {error_msg}"

            error_msg += f"\n\n期望包含:\n{normalize_code(expected)}\n\n实际代码:\n{normalize_code(code)}"
            raise AssertionError(error_msg)


def compile_source(
    obj_or_src: Union[Callable, str],
    module_name: str = "main",
    outdir: Path = Path("."),
) -> tuple[str, str]:
    if isinstance(obj_or_src, Callable):
        src = textwrap.dedent(inspect.getsource(obj_or_src))
        func_name: Optional[str] = obj_or_src.__name__
    else:
        src = str(obj_or_src)
        func_name = None

    # 获取 stubs 目录
    import os

    stubs_dir = Path(__file__).parent.parent.parent.parent / "stubs"

    compiler = MCPyCompiler(
        {
            "compiler": {
                "stubs_directories": [str(stubs_dir)] if stubs_dir.exists() else []
            }
        }
    )
    compiler.config.update(
        {
            "output_dir": str(outdir),
            "verbose": False,
            "format_code": True,
            "fail_on_checks": False,
            "source_root": ".",
            "save_generated": False,
            "project_name": module_name,
        }
    )
    res = compiler.compile_source(src, func_name or "snippet.py")
    if not res.success:
        raise AssertionError(f"codegen failed: {res.errors}")
    if not res.generated_contents:
        return ("", "")
    return (
        res.generated_contents[f"{module_name}.h"],
        res.generated_contents[f"{module_name}.cpp"],
    )


# 合并所有生成的文件内容
def merged_contents(contents, module_name):
    hpp = f"// @mcpy_file {module_name}.h\n\n{contents[0]}"
    cpp = f"// @mcpy_file {module_name}.cpp\n\n{contents[1]}"
    return hpp + "\n\n" + cpp


def _extract_function_block(cpp: str, func_name: str) -> str:
    """更健壮的函数块提取算法

    算法步骤：
    1. 找到函数名位置
    2. 向前查找返回类型
    3. 向后匹配参数列表（处理嵌套括号）
    4. 匹配函数体（处理嵌套大括号）
    """
    # 查找所有函数名出现的位置
    start_pos = 0
    while True:
        func_pos = cpp.find(func_name, start_pos)
        if func_pos == -1:
            break

        # 检查函数名后是否有左括号（跳过空白、换行和注释）
        pos = func_pos + len(func_name)
        found_paren = False

        while pos < len(cpp):
            char = cpp[pos]
            if char.isspace():
                pos += 1
            elif char == "(":
                found_paren = True
                break
            elif char == "/" and pos + 1 < len(cpp):
                if cpp[pos : pos + 2] == "//":
                    # 单行注释，跳到下一行
                    pos = cpp.find("\n", pos)
                    if pos == -1:
                        break
                    pos += 1
                elif cpp[pos : pos + 2] == "/*":
                    # 多行注释，跳过到 */
                    pos = cpp.find("*/", pos + 2)
                    if pos == -1:
                        break
                    pos += 2
                else:
                    # 不是注释，说明不是函数定义
                    break
            else:
                # 遇到其他非空白、非括号字符，说明不是函数定义
                break

        if found_paren:
            # 找到了函数定义，提取完整函数
            result = _extract_complete_function(cpp, func_pos, func_name)
            if result is not None:
                return result
            # 如果是函数声明，继续查找下一个

        start_pos = func_pos + 1

    return ""


def _extract_complete_function(cpp: str, func_pos: int, func_name: str) -> str | None:
    """从指定位置提取完整函数"""
    # 1. 向前查找返回类型
    return_type = _extract_return_type(cpp, func_pos)

    # 2. 找到函数名后的左括号位置
    paren_pos = func_pos + len(func_name)
    while paren_pos < len(cpp) and cpp[paren_pos] != "(":
        paren_pos += 1

    if paren_pos >= len(cpp):
        return return_type + " " + func_name if return_type else func_name

    # 3. 匹配参数列表的括号
    params_end = _find_matching_bracket(cpp, paren_pos, "(", ")")
    if params_end == -1:
        return cpp[func_pos - len(return_type) - 1 if return_type else func_pos :]

    # 4. 找到函数体开始的大括号
    brace_pos = params_end + 1
    while brace_pos < len(cpp) and cpp[brace_pos] != "{":
        if cpp[brace_pos] == ";":  # 函数声明，不是定义，返回None表示继续查找
            return None
        brace_pos += 1

    if brace_pos >= len(cpp):
        return cpp[
            func_pos - len(return_type) - 1 if return_type else func_pos : params_end
            + 1
        ]

    # 5. 匹配函数体的大括号
    body_end = _find_matching_bracket(cpp, brace_pos, "{", "}")
    if body_end == -1:
        return cpp[func_pos - len(return_type) - 1 if return_type else func_pos :]

    # 6. 提取完整函数
    start_pos = func_pos - len(return_type) - 1 if return_type else func_pos
    return cpp[start_pos : body_end + 1].strip()


def _extract_return_type(cpp: str, func_pos: int) -> str:
    """提取函数返回类型"""
    # 向前查找到上一个分号/大括号或行首
    start_pos = func_pos - 1
    line_start = func_pos

    # 先找到当前行的开始位置
    while line_start > 0 and cpp[line_start - 1] != "\n":
        line_start -= 1

    while start_pos >= line_start:  # 限制在当前行内
        char = cpp[start_pos]
        if char in ";{}":
            start_pos += 1
            break
        start_pos -= 1
    else:
        # 如果当前行没有找到分号/大括号，向前查找但限制范围
        while start_pos >= 0:
            char = cpp[start_pos]
            if char in ";{}":
                start_pos += 1
                break
            elif char == "\n":
                # 检查前一行是否是预处理指令或其他非函数内容
                prev_line_start = start_pos
                while prev_line_start > 0 and cpp[prev_line_start - 1] != "\n":
                    prev_line_start -= 1
                prev_line = cpp[prev_line_start:start_pos].strip()
                if (
                    prev_line.startswith("#")
                    or prev_line.startswith("//")
                    or prev_line == ""
                ):
                    # 遇到预处理指令、注释或空行，停止查找
                    start_pos += 1
                    break
            start_pos -= 1
        else:
            start_pos = 0

    # 提取返回类型部分，去除前后空白和换行
    return_type = cpp[start_pos:func_pos].strip()
    return return_type


def _find_matching_bracket(
    cpp: str, start_pos: int, open_bracket: str, close_bracket: str
) -> int:
    """找到匹配的括号位置"""
    bracket_count = 1
    pos = start_pos + 1
    in_string = False
    string_char = None
    escape_next = False

    while pos < len(cpp) and bracket_count > 0:
        char = cpp[pos]

        # 处理字符串和转义
        if escape_next:
            escape_next = False
        elif char == "\\" and in_string:
            escape_next = True
        elif char in ['"', "'"] and not escape_next:
            if not in_string:
                in_string = True
                string_char = char
            elif char == string_char:
                in_string = False
                string_char = None

        # 在字符串外才计算括号
        if not in_string:
            if char == open_bracket:
                bracket_count += 1
            elif char == close_bracket:
                bracket_count -= 1

        pos += 1

    return pos - 1 if bracket_count == 0 else -1


def compile_chunk(obj_or_src: Union[Callable, str], module_name: str = "main") -> str:
    """编译函数对象或源码片段为 C++，返回该函数的实现块文本。

    - 传入函数对象（可为内嵌函数）: 提取源码，去缩进后作为模块级函数进行编译，并返回该函数实现块。
    - 传入源码字符串: 作为完整模块源码进行编译，返回完整 C++ 源文件内容。
    """
    func_name = obj_or_src.__name__ if isinstance(obj_or_src, Callable) else None
    contents = compile_source(obj_or_src, module_name)
    txt = merged_contents(contents, module_name)
    return _extract_function_block(txt, func_name) if func_name else txt


def compile_file(
    path: Union[str, Path], func_name: Optional[str] = None, module_name: str = "main"
) -> str:
    """编译文件路径为 C++，可选指定函数名仅返回该函数实现块。

    - path: 源文件路径
    - func_name: 若提供，则仅提取并返回该函数实现块；否则返回完整 C++ 源文件内容
    """
    p = Path(path)
    src = p.read_text(encoding="utf-8")
    contents = compile_source(src, module_name)
    txt = merged_contents(contents, module_name)
    return _extract_function_block(txt, func_name) if func_name else txt
